#!/usr/bin/python

import json
import sys
import numpy as np
import matplotlib.pyplot as plt
import collections
import ast
import pickle
import multiprocessing
import os
from time import sleep
from time import time
from itertools import groupby
from sklearn.pipeline import Pipeline
from sklearn.naive_bayes import MultinomialNB
from sklearn.feature_extraction.text import TfidfTransformer
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.metrics import confusion_matrix
from sklearn.datasets import make_multilabel_classification
from sklearn.preprocessing import MultiLabelBinarizer
from sklearn.multiclass import OneVsRestClassifier
from sklearn.svm import LinearSVC
from sklearn.multiclass import OneVsOneClassifier
from sklearn import datasets
from sklearn.externals import joblib



validLabels = None
StatFile = None
CMFile = None
PHFile = None
PHTable = None
LabelDict = None

def main():
	
	# Done
	#train(1,'PER_utf8','Label_L1_last',ScikitMultinomial)
	#train(3,'PER_utf8','Label_L1_last',ScikitMultinomial)
	#train(9,'PER_utf8','Label_L1_last',ScikitMultinomial)
	#train(1,'PER_ascii','Label_L1_last',ScikitMultinomial)
	#train(3,'PER_ascii','Label_L1_last',ScikitMultinomial)
	#train(9,'PER_ascii','Label_L1_last',ScikitMultinomial)
	#train(1,'PER_ascii_stem','Label_L1_last',ScikitMultinomial)
	#train(3,'PER_ascii_stem','Label_L1_last',ScikitMultinomial)
	#train(9,'PER_ascii_stem','Label_L1_last',ScikitMultinomial)
	#train(1,'PER_ascii_stem_stopword','Label_L1_last',ScikitMultinomial)
	#train(3,'PER_ascii_stem_stopword','Label_L1_last',ScikitMultinomial)
	#train(9,'PER_ascii_stem_stopword','Label_L1_last',ScikitMultinomial)
	#train(9,'PER_ascii_stem_stopword','Label_L1_first',ScikitMultinomial)
	#train(9,'PER_ascii_stem','Label_L2_last',ScikitMultinomial)
	#train(1,'PER_ascii_stem_stopword','Label_L2_last',ScikitMultinomial)
	#train(9,'PER_ascii_stem_stopword','Label_L2_last',ScikitMultinomial)
	
	#train(1,'PER_ascii_stem_stopword','Label_Leaves',ML_ScikitMultinomial)
	#train(3,'PER_ascii_stem_stopword','Label_Leaves',ML_ScikitMultinomial)
	#train(9,'PER_ascii_stem_stopword','Label_Leaves',ML_ScikitMultinomial)

	#train(9,'PER_ascii_stem_stopword','Label_Leaves',NB)
	
	# Not Run Yet
	train(9,'PER_ascii_stem_stopword','Label_Leaves',ML_NB)
	

	
	#train(1,'PER_ascii_stem_stopword','Label_L1_last',ScikitOvR)
	


def train(trainThres,data,label,model):
	modelName=model.__name__
	#cnt= 0 - 337060

	#trainThres = 1/3/9
	#33706 1
	#101118 3
	#303354 9

	#train cnt % 10 < trainThres
	#test cnt % 10 == 9
	
	print '--------------------------'
	print 'trainThres = ',trainThres
	print 'modelName = ',modelName
	print 'data = ',data
	print 'label = ',label

	dataFile = open('../'+data,'r')
	labelFile = open('../'+label,'r')

	#StatFile = open('../output/Stats_'+modelName+'_'+str(trainThres)+'_'+data+'_'+label+'.txt','w')
	
	global PHFile, PHTable
	PHFile = open('../ParentHash.json','r')
	PHTable = json.load(PHFile)

	global LabelDict	
	LabelDict = dict()
	for l in range(1,10):
		LabelDict[l] = np.loadtxt('../validLabels_L'+str(l)+'.txt',dtype=str)

	global validLabels
	validLabels = np.loadtxt('../validLabels_L1.txt',dtype=str)

	trainData = []
	trainLabel = []
	
	cnt=0
	for line in dataFile:
		#print 'Read Data #',cnt+1
		if (cnt % 10) == 9:
			pass
		elif (cnt % 10) < trainThres:
			trainData.append(line.strip())
		cnt += 1
	print 'Finish Reading Data'

	cnt=0
	for line in labelFile:
		#print 'Read Label #',cnt+1
		if (cnt % 10) == 9:
			pass
		elif (cnt % 10) < trainThres:
			trainLabel.append(line.strip())
		cnt += 1
	print 'Finish Reading Labels'

	print 'Train Data Size = ',len(trainData)
	print 'Train Label Size = ',len(trainLabel)
	print ''
	
	model(trainThres,trainData,trainLabel)

def ML_NB_thread(trainThres, cls, trainData, trainLabels, clf):

	subData = []
	subLabels = []

	for j in range(len(trainLabels)):
		ret,label = hasAncestor(trainLabels[j],clf)
		if ret:
			if not trainData[j] == '':
				subData.append(trainData[j])
				subLabels.append(label)

	#print 'Data Size = ',len(subLabels)
	if len(subLabels) <= 5:
		return

	cls = cls.fit(subData, subLabels)
	
	path = '../../../Data229/CLF/ML_NB'+str(trainThres)+'/'
	if not os.path.exists(path):
		os.makedirs(path)
	joblib.dump(cls, path+clf+'.pkl')
	
def ML_NB(trainThres,trainData,trainLabels):
	#Train
	time0 = time()
	main_thread = multiprocessing.current_process()

	text_clf = dict()
	text_clf['0'] = Pipeline([('vect', CountVectorizer()),
						('tfidf', TfidfTransformer()),
						('clf', MultinomialNB()),])

	for i in trainLabels:
		curr = i
		while curr != '0' :
			if (PHTable[curr] not in text_clf): # and (PHTable[curr] in LabelDict[1]) and False:
				text_clf[PHTable[curr]]=Pipeline([('vect', CountVectorizer()),
						('tfidf', TfidfTransformer()),
						('clf', MultinomialNB()),])
			curr = PHTable[curr]
	print 'CLFs allocated'
	CLFsize = len(text_clf)
	print 'CLFs Size = ',CLFsize

	cnt=0
	for clf in text_clf:

		cnt+=1

		#print '-'
		print 'Training CLF',cnt,'/',CLFsize, '(',clf,')','time =',(time()-time0)


		while True:
			if  len(multiprocessing.active_children()) < 8:
				multiprocessing.Process(target=ML_NB_thread, args=(trainThres,text_clf[clf],trainData,trainLabels,clf)).start()
				break
			else:
				sleep(0.1)

	while len(multiprocessing.active_children()) != 0:
		print 'wait for subprocess to finish'
		sleep(1)
		
	print 'Training is Done'
	
def covers(c,p):
	
	if c == p:
		return True
	elif c == 0:
		return False
	else:
		return covers(PHTable[str(c)],p)

def hasAncestor(curr,p):
	#print 'c=',type(curr),curr
	#print 'p=',type(p),p

	if curr == '0':
		return False,curr

	if PHTable[curr] == p:
		return True,curr
	else:
		return hasAncestor(PHTable[curr],p)
	#else:
		#print 'F'
		#exit(0)
	#	return False,curr

	
if __name__ == "__main__":
    main()
